package ar.edu.itba.it.raytracer.engine

import Vector._
import RichVector._
import ar.edu.itba.it.raytracer.engine.Color._
import ar.edu.itba.it.raytracer.scene.Geom.eps
import ar.edu.itba.it.raytracer.scene.{ AABB, Camera, Geom, Material, Point, SceneObject, Plane, Sphere }
import java.awt.image.{ BufferedImage, ImageObserver }
import java.awt.{ Dimension, Graphics2D, Image }
import java.io.File
import javax.imageio.ImageIO
import scala.Math.{ max, pow }
import scala.swing.event.MouseMoved
import scala.swing.{ MainFrame, Panel, SimpleSwingApplication }

object RayTracer extends SimpleSwingApplication {

  private val rayTraceCount = 3

  private val (camera, scene) = createScene4

  def top = new MainFrame {
    title = "Ray Tracer"
    val image = new RayTracer(camera, scene).compute
    ImageIO.write(image, "jpeg", File.createTempFile("temp", "jpeg"));

    preferredSize = new Dimension(camera.xres, camera.yres)
    contents = new Panel() {
      listenTo(mouse.moves)
      reactions += {
        case MouseMoved(src, point, modifiers) => println(point)
      }

      override def paint(g: Graphics2D) {
        super.paint(g)
        g.drawImage(image, 0, 0, new ImageObserver() {
          override def imageUpdate(img: Image, arg2: Int, arg3: Int, arg4: Int, arg5: Int, arg6: Int) = true
        })
      }
    }
  }

  lazy val createScene1 = (new Camera(Vector(-8, 0, 0), Vector(0, 1, 0), Vector(4, 0, 0), 14, 10, 700, 500), List(
    SceneObject.create(Plane(Vector(0, -3, 0), Vector(0, 1, 0)), Material(Color(0.4, 0.3, 0.3), 0, 1, 0)),
    SceneObject.create(Sphere(Vector(10, 0, -4.5), 2), Material(blue, 0.6, 1, 0.5)),
    SceneObject.create(Sphere(Vector(10, 0, 0), 2), Material(green, 0.3, 0.2, 0.7)),
    SceneObject.create(Sphere(Vector(10, 0, 4.5), 2), Material(red, 0.2, 0.8, 0.3)),
    SceneObject.light(AABB(Vector(0, 20, 0), 2), white)))
  //    SceneObject.light(Point(Vector(0, 50, 2)), white)))

  lazy val createScene = (new Camera(Vector(-8, 0, 0), Vector(0, 1, 0), Vector(7, 0, 0), 14, 10, 350, 250), List(
    SceneObject.create(Plane(Vector(0, -3, 0), Vector(0, 1, 0)), Material(Color(0.4, 0.3, 0.3), 0, 1, 0)),
    SceneObject.create(Sphere(Vector(10, 0, 0), 2), Material(green, 0.3, 1, 0.7)),
    SceneObject.light(AABB(Vector(10, 5, 0), 1), white)))
  //    SceneObject.light(Point(Vector(5, 50, 0)), white)))

  lazy val createScene2 = (new Camera(Vector(0, 0, -5), Vector(0, 1, 0), Vector(0, 0, 0), 8, 6, 800, 600), List(
    SceneObject.create(Plane(Vector(0, -4, 0), Vector(0, 1, 0)), Material(Color(0.4, 0.3, 0.3), 0, 1, 0)),
    SceneObject.create(Sphere(Vector(-2.5, -0.8, 3), 2.5), Material(Color(0.7, 0.7, 0.7), 0.6, 1, 1)),
    SceneObject.create(Sphere(Vector(2, -0.5, 7), 2), Material(Color(0.7, 0.7, 1), 1, 0.1, 0.5)),
    SceneObject.light(Point(Vector(2, 5, 1)), Color(0.6, 0.6, 0.8)),
    SceneObject.light(Point(Vector(0, 5, 5)), Color(0.4, 0.4, 0.4))))

  lazy val createScene3 = (new Camera(Vector(6, 0, -8), Vector(0, 1, 0), Vector(0, 0, 0), 8, 6, 200, 150), List(
    SceneObject.create(AABB(Vector(-1.5, -1.5, 2), 3), Material(red, 0, 1, 0)),
    SceneObject.light(Point(Vector(10, 0, 0)), white),
    SceneObject.light(Point(Vector(0, 20, -10)), white)))

  lazy val createScene4 = (new Camera(Vector(0, 4, -10), Vector(0, 1, 0), Vector(0, 3.5, 0), 10, 7, 500, 350), List(
    SceneObject.create(Plane(Vector(0, 0, 0), Vector(0, 1, 0)), Material(Color(0.4, 0.3, 0.3), 0, 1, 0.1)),
    SceneObject.create(Plane(Vector(0, 0.5, 10), Vector(0, 0, -1)), Material(Color(0.4, 0.3, 0.3), 0, 1, 0.2)),
    SceneObject.create(Sphere(Vector(0, 4, 7), 2), Material(Color(0.7, 0.7, 1), 0.2, 0.3, 0.5)),
    SceneObject.create(Sphere(Vector(-5, 4, 7), 2), Material(Color(0.7, 0.7, 1), 0.5, 0.6, 0.2)),
    SceneObject.create(Sphere(Vector(5, 4, 7), 2), Material(Color(0.7, 0.7, 1), 0.5, 0.6, 0.4)),
//    SceneObject.light(AABB(Vector(0, 9, 4), 1), white),
    SceneObject.light(AABB(Vector(0, 9, 7), 1), white)))
//    SceneObject.light(Point(Vector(0, 9, 8)), white)))
}

class RayTracer(camera: Camera, scene: List[SceneObject]) {
  import RayTracer._

  val lights = scene.filter(_.light)

  def compute = {
    val width = camera.xres
    val height = camera.yres
    val result: Array[Array[Color]] = Array.tabulate(camera.yres, camera.xres) { (x, y) => black }
    val pixels = for (y <- 0 to camera.yres - 1; x <- 0 to camera.xres - 1) yield ((x, y))
    pixels.par.foreach({ case (x, y) => result(y)(x) = color(camera.ray(y, x), rayTraceCount) })

    val image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB)
    for (i <- 0 to camera.yres - 1; j <- 0 to camera.xres - 1) {
      val color = result(i)(j)
      image.getRaster.setPixel(j, i, Array((color.red * 255).toInt, (color.green * 255).toInt, (color.blue * 255).toInt))
    }
    image
  }

  private def color(ray: Ray, traceLeft: Int): Color = {
    var col = Color.black
    val intersection = intersect(ray)
    if (intersection.isDefined) {
      val (obj, poi, t) = intersection.get
      val normal = obj.geom.normal(poi, ray)
      val r = ray.direction - 2 * (ray.direction * normal) * normal
      if (traceLeft - 1 >= 1 && obj.material.reflection > 0) col += color(Ray(poi, r), traceLeft - 1) * obj.material.reflection

      for (light <- lights) {
        val (shade, dirToLight) = computeShade(poi, light)
        if (shade > 0) {
          val diffuse = diffuseLight(light, obj, dirToLight, normal)
          if (diffuse.isDefined) col += diffuse.get * shade

          val specular = specularLight(obj, ray, light, dirToLight, normal)
          if (specular.isDefined) col += specular.get * shade
        }
      }
    }
    col
  }

  private def computeShade(poi: Vector, light: SceneObject): (Double, Vector) = {
    def rays(aabb: AABB) = for (i <- 0 to 7; j <- 0 to 7) yield (Ray(poi, (Vector(aabb.v1.x + aabb.size * i / 8, aabb.v1.y, aabb.v1.z + aabb.size * j / 8) - poi).normalized))
    light.geom match {
      case aabb: AABB => {
        val dirToLight = ((aabb.v1 + Vector(0.5 * aabb.size, 0.5 * aabb.size, 0.5 * aabb.size)) - poi).normalized
        (rays(aabb).count(intersect(_, light)) / 64.0, dirToLight)
      }
      case point: Point => {
        val dirToLight = (point.position - poi).normalized
        val rayToLight = Ray(poi + dirToLight * eps, dirToLight)
        val intersection = intersect(rayToLight).get
        val shade = if (intersection._1 == light) 1 else 0
        (shade, dirToLight)
      }
    }
  }

  private def diffuseLight(light: SceneObject, obj: SceneObject, directionToLight: Vector, normal: Vector) = {
    var result: Option[Color] = None
    val dot = directionToLight * normal
    if (dot > 0) result = Some(light.material.color * obj.material.color * dot * obj.material.diffuse)
    result
  }

  private def specularLight(obj: SceneObject, ray: Ray, light: SceneObject, directionToLight: Vector, normal: Vector) = {
    var result: Option[Color] = None
    val r = directionToLight - 2 * (directionToLight * normal) * normal
    val dot = ray.direction * r
    if (dot > 0) { result = Some(light.material.color * pow(dot, 5) * obj.material.specular) }
    result
  }

  private def intersect(ray: Ray): Option[(SceneObject, Vector, Double)] = {
    val objects = for {
      obj <- scene
      poi = obj.geom.intersect(ray)
      if (poi.isDefined)
      t = ray.tFor(poi.get)
      if (t.isDefined && t.get > eps)
    } yield (obj, poi.get, t.get)
    if (objects.isEmpty) None else Some(objects.sortBy({ case (obj, poi, t) => t }).apply(0))
  }

  private def intersect(ray: Ray, obj: SceneObject): Boolean = intersect(ray).map(_._1 == obj).getOrElse(false)
}